{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c44a8591",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers.models.embeddings import TimestepEmbedding, get_timestep_embedding\n",
    "from diffusers.models.unet_2d_blocks import CrossAttnDownBlock2D, DownBlock2D\n",
    "from diffusers.models.unet_2d_blocks import UNetMidBlock2DCrossAttn\n",
    "from transformers import PreTrainedModel, PretrainedConfig\n",
    "\n",
    "#计算unet模型的down和mid部分的数据\n",
    "class ControlNet(PreTrainedModel):\n",
    "    config_class = PretrainedConfig\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "\n",
    "        #入参embed部分\n",
    "        self.out_vae_noise_embed = torch.nn.Conv2d(4,\n",
    "                                                   320,\n",
    "                                                   kernel_size=3,\n",
    "                                                   padding=1)\n",
    "\n",
    "        self.time_embed = TimestepEmbedding(320, 1280, act_fn='silu')\n",
    "\n",
    "        self.condition_embed = torch.nn.Sequential(\n",
    "            torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(32, 96, kernel_size=3, stride=2, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(96, 256, kernel_size=3, stride=2, padding=1),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(256, 320, kernel_size=3, stride=1, padding=1),\n",
    "        )\n",
    "\n",
    "        #unet的down部分\n",
    "        self.unet_down = torch.nn.ModuleList([])\n",
    "        for i in range(3):\n",
    "            self.unet_down.append(\n",
    "                CrossAttnDownBlock2D(num_layers=2,\n",
    "                                     in_channels=[320, 320, 640][i],\n",
    "                                     out_channels=[320, 640, 1280][i],\n",
    "                                     temb_channels=1280,\n",
    "                                     add_downsample=True,\n",
    "                                     resnet_eps=1e-5,\n",
    "                                     resnet_act_fn='silu',\n",
    "                                     resnet_groups=32,\n",
    "                                     downsample_padding=1,\n",
    "                                     cross_attention_dim=768,\n",
    "                                     attn_num_head_channels=8,\n",
    "                                     dual_cross_attention=False,\n",
    "                                     use_linear_projection=False,\n",
    "                                     only_cross_attention=False,\n",
    "                                     upcast_attention=False,\n",
    "                                     resnet_time_scale_shift='default'))\n",
    "        self.unet_down.append(\n",
    "            DownBlock2D(num_layers=2,\n",
    "                        in_channels=1280,\n",
    "                        out_channels=1280,\n",
    "                        temb_channels=1280,\n",
    "                        add_downsample=False,\n",
    "                        resnet_eps=1e-5,\n",
    "                        resnet_act_fn='silu',\n",
    "                        resnet_groups=32,\n",
    "                        downsample_padding=1,\n",
    "                        resnet_time_scale_shift='default'))\n",
    "\n",
    "        #unet的mid部分\n",
    "        self.unet_mid = UNetMidBlock2DCrossAttn(\n",
    "            in_channels=1280,\n",
    "            temb_channels=1280,\n",
    "            resnet_eps=1e-5,\n",
    "            resnet_act_fn='silu',\n",
    "            output_scale_factor=1,\n",
    "            resnet_time_scale_shift='default',\n",
    "            cross_attention_dim=768,\n",
    "            attn_num_head_channels=8,\n",
    "            resnet_groups=32,\n",
    "            use_linear_projection=False,\n",
    "            upcast_attention=False)\n",
    "\n",
    "        #control的down部分\n",
    "        self.control_down = torch.nn.ModuleList([\n",
    "            torch.nn.Conv2d(320, 320, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(320, 320, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(320, 320, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(320, 320, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(640, 640, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(640, 640, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(640, 640, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(1280, 1280, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(1280, 1280, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(1280, 1280, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(1280, 1280, kernel_size=1, stride=1, padding=0),\n",
    "            torch.nn.Conv2d(1280, 1280, kernel_size=1, stride=1, padding=0),\n",
    "        ])\n",
    "\n",
    "        #control的mid部分\n",
    "        self.control_mid = torch.nn.Conv2d(1280, 1280, kernel_size=1)\n",
    "\n",
    "    def forward(self, out_vae_noise, noise_step, out_encoder, condition):\n",
    "        #out_vae_noise -> [1, 4, 64, 64]\n",
    "        #noise_step -> [1]\n",
    "        #out_encoder -> [1, 77 ,768]\n",
    "        #condition -> [1, 3, 512, 512]\n",
    "\n",
    "        #编码noise_step\n",
    "        #[1] -> [1, 320]\n",
    "        noise_step = get_timestep_embedding(noise_step,\n",
    "                                            320,\n",
    "                                            flip_sin_to_cos=True,\n",
    "                                            downscale_freq_shift=0)\n",
    "        #[1, 320] -> [1, 1280]\n",
    "        noise_step = self.time_embed(noise_step, None)\n",
    "\n",
    "        #out_vae_noise升到高维\n",
    "        #[1, 4, 64, 64] -> [1, 320, 64, 64]\n",
    "        out_vae_noise = self.out_vae_noise_embed(out_vae_noise)\n",
    "\n",
    "        #condition投影到和out_vae_noise同一维度空间\n",
    "        #[1, 3, 512, 512] -> [1, 320, 64, 64]\n",
    "        condition = self.condition_embed(condition)\n",
    "\n",
    "        #向out_vae_noise中添加condition信息\n",
    "        #[1, 320, 64, 64]\n",
    "        out_vae_noise += condition\n",
    "\n",
    "        #unet的down部分计算,每一层当中包括了3个串行的注意力计算,所以每一层都有3个计算结果.\n",
    "        #[1, 320, 64, 64]\n",
    "        #[1, 320, 64, 64]\n",
    "        #[1, 320, 64, 64]\n",
    "        #[1, 320, 32, 32]\n",
    "        #[1, 640, 32, 32]\n",
    "        #[1, 640, 32, 32]\n",
    "        #[1, 640, 16, 16]\n",
    "        #[1, 1280, 16, 16]\n",
    "        #[1, 1280, 16, 16]\n",
    "        #[1, 1280, 8, 8]\n",
    "        #[1, 1280, 8, 8]\n",
    "        #[1, 1280, 8, 8]\n",
    "        out_unet_down = [out_vae_noise]\n",
    "        for i in range(4):\n",
    "            if i < 3:\n",
    "                #这里只记录了out_vae_noise的维度变换,输出的维度看上面的out_unet_down\n",
    "                #[1, 320, 64, 64] -> [1, 320, 32, 32]\n",
    "                #[1, 320, 32, 32] -> [1, 640, 16, 16]\n",
    "                #[1, 640, 16, 16] -> [1, 1280, 8, 8]\n",
    "                out_vae_noise, out = self.unet_down[i](\n",
    "                    hidden_states=out_vae_noise,\n",
    "                    temb=noise_step,\n",
    "                    encoder_hidden_states=out_encoder,\n",
    "                    attention_mask=None,\n",
    "                    cross_attention_kwargs=None)\n",
    "            else:\n",
    "                #[1, 1280, 8, 8] -> [1, 1280, 8, 8]\n",
    "                out_vae_noise, out = self.unet_down[i](\n",
    "                    hidden_states=out_vae_noise, temb=noise_step)\n",
    "\n",
    "            out_unet_down.extend(out)\n",
    "\n",
    "        #unet的mid计算,维度不变\n",
    "        #[1, 1280, 8, 8] -> [1, 1280, 8, 8]\n",
    "        out_vae_noise = self.unet_mid(out_vae_noise,\n",
    "                                      noise_step,\n",
    "                                      encoder_hidden_states=out_encoder,\n",
    "                                      attention_mask=None,\n",
    "                                      cross_attention_kwargs=None)\n",
    "\n",
    "        #control的down的部分计算,维度不变,两两组合,分别计算即可\n",
    "        out_control_down = [\n",
    "            self.control_down[i](out_unet_down[i]) for i in range(12)\n",
    "        ]\n",
    "\n",
    "        #control的mid的部分计算,维度不变\n",
    "        out_control_mid = self.control_mid(out_vae_noise)\n",
    "\n",
    "        return out_control_down, out_control_mid\n",
    "\n",
    "\n",
    "# controlnet = ControlNet(PretrainedConfig())\n",
    "\n",
    "# out_control_down, out_control_mid = controlnet(\n",
    "#     out_vae_noise=torch.randn(1, 4, 64, 64),\n",
    "#     noise_step=torch.LongTensor([400]),\n",
    "#     out_encoder=torch.randn(1, 77, 768),\n",
    "#     condition=torch.randn(1, 3, 512, 512))\n",
    "\n",
    "# for i in out_control_down:\n",
    "#     print(i.shape)\n",
    "\n",
    "# out_control_mid.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e3f084c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#从unet当中抽取参数\n",
    "def load_params(controlnet, unet):\n",
    "    controlnet.out_vae_noise_embed.load_state_dict(unet.conv_in.state_dict())\n",
    "    controlnet.time_embed.load_state_dict(unet.time_embedding.state_dict())\n",
    "    controlnet.unet_down.load_state_dict(unet.down_blocks.state_dict())\n",
    "    controlnet.unet_mid.load_state_dict(unet.mid_block.state_dict())\n",
    "\n",
    "\n",
    "# from diffusers import UNet2DConditionModel\n",
    "\n",
    "# unet = UNet2DConditionModel.from_pretrained('runwayml/stable-diffusion-v1-5',\n",
    "#                                             subfolder='unet')\n",
    "\n",
    "# load_params(controlnet, unet)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
